"""
六足机器人运动控制器
实现六足机器人的特定运动控制功能
"""

import math
from typing import Dict, Any, List, Tuple
from dataclasses import dataclass
from .base_controller import BaseMotionController
from sensors.imu_sensor import IMUSensor
import utils.filter as filter


@dataclass
class LegState:
    """腿部状态"""

    position: Tuple[float, float, float]  # 末端位置 (x, y, z)
    velocity: Tuple[float, float, float]  # 末端速度
    force: Tuple[float, float, float]  # 末端受力
    contact: bool  # 是否接触地面


class HexapodController(BaseMotionController):
    """六足机器人运动控制器"""

    # 运动学参数
    C = 44.60  # 髋关节到膝关节的距离
    F = 75.00  # 膝关节到踝关节的距离
    T = 126.50  # 踝关节到足端的距离
    ANGLE_FACTOR = 180 / math.pi / 0.24  # 角度转换因子

    # 预定义姿态参数
    STAND_POSITION = [100.0, 100.0, -70.0]  # 站立姿态
    SIT_POSITION = [100.0, 100.0, 20.0]  # 坐下姿态
    LIFT_POSITION = [100.0, 100.0, -40.0]  # 抬起姿态

    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.imu = IMUSensor(0x79)  # 初始化IMU传感器
        self.imu.set_gyro_range(self.imu.GYRO_RANGE_2000DEG)
        self.leg_states = [
            LegState((0, 0, 0), (0, 0, 0), (0, 0, 0), False) for _ in range(6)
        ]
        self._init_camera()

    def _init_camera(self):
        """初始化摄像头位置"""
        self.set_servo_position(1, 1500, 100)  # 水平舵机
        self.set_servo_position(2, 1500, 100)  # 垂直舵机

    async def stand(self) -> None:
        """设置机器人站立姿态"""
        for leg in range(6):
            await self.move_leg(leg, self.STAND_POSITION, 1000)

    async def sit(self) -> None:
        """设置机器人坐下姿态"""
        for leg in range(6):
            await self.move_leg(leg, self.SIT_POSITION, 1000)

    async def reset(self) -> None:
        """重置机器人姿态"""
        await self.stand()
        self._init_camera()

    async def lift_legs(self, legs: List[int]) -> None:
        """抬起指定的腿"""
        for leg in legs:
            if 0 <= leg < 6:
                await self.move_leg(leg, self.LIFT_POSITION, 500)

    async def lower_legs(self, legs: List[int]) -> None:
        """放下指定的腿"""
        for leg in legs:
            if 0 <= leg < 6:
                await self.move_leg(leg, self.STAND_POSITION, 500)

    async def move_leg(self, leg: int, position: List[float], speed: int) -> None:
        """控制单条腿运动"""
        angles = self.get_leg_angles(leg, position, speed)
        for i in range(3):
            await self.set_servo_position(leg * 3 + i + 1, int(angles[i]), speed)

    def get_leg_angles(
        self, leg: int, position: List[float], speed: int
    ) -> List[float]:
        """计算腿部关节角度"""
        angles = []
        output = []

        # 计算髋关节角度
        angle_hip = math.atan(position[1] / position[0])
        angles.append(angle_hip)

        # 计算腿部长度
        L = position[1] / math.sin(angle_hip)
        temp = math.pow(position[2], 2) + pow(L - self.C, 2)
        ft = math.sqrt(temp)

        # 计算膝关节角度
        a = math.atan(position[2] / (L - self.C))
        b = math.acos(
            (math.pow(self.F, 2) + math.pow(ft, 2) - math.pow(self.T, 2))
            / (2 * self.F * ft)
        )
        angles.append(a + b)

        # 计算踝关节角度
        angles.append(
            math.acos(
                (math.pow(ft, 2) - math.pow(self.F, 2) - math.pow(self.T, 2))
                / (2 * self.F * self.T)
            )
        )

        # 转换为舵机角度
        if leg < 3:
            output.append(313 + angles[0] * self.ANGLE_FACTOR)
            output.append(500 - angles[1] * self.ANGLE_FACTOR)
            output.append(687 - angles[2] * self.ANGLE_FACTOR - 5)
        else:
            output.append(687 - angles[0] * self.ANGLE_FACTOR)
            output.append(500 + angles[1] * self.ANGLE_FACTOR)
            output.append(313 + angles[2] * self.ANGLE_FACTOR + 5)

        return output

    async def get_imu_data(self) -> Dict[str, Any]:
        """获取IMU数据"""
        accel = self.imu.get_accel_data()
        gyro = self.imu.get_gyro_data()
        temp = self.imu.get_temp()

        # 使用互补滤波计算姿态
        roll = filter.filter_r(accel["x"], gyro["y"])
        pitch = filter.filter_p(accel["y"], gyro["x"])

        return {
            "accel": accel,
            "gyro": gyro,
            "temp": temp,
            "roll": roll,
            "pitch": pitch,
        }

    def close(self):
        """关闭资源"""
        super().close()
        # 关闭IMU
        if hasattr(self, "imu"):
            self.imu.close()
